package com.bsfcs.rag.service;

import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;

import java.util.ArrayList;
import java.util.List;


public class ChineseTextSplitter {

    private TokenTextSplitter tokenSplitter;

    public ChineseTextSplitter() {
        // 初始化 TokenTextSplitter
        tokenSplitter = new TokenTextSplitter(
                350,    // 每块目标 token 大小
                100,    // 块之间最小重叠 token
                5,      // 切分尝试次数
                1,      // 块大小容忍度
                true    // 保持标题在块开头
        );
    }

    /**
     * 中文文档切分
     * @param text 原始中文文本
     * @return 切分后的块列表
     */
    public List<Document> splitChineseDocument(String text) {
        if (text == null || text.isEmpty()) {
            return new ArrayList<>();
        }
        String[] paragraphs = text.split("\n\n+");
        List<Document> initialSegments = new ArrayList<>();

        for (String para : paragraphs) {
            if (!para.trim().isEmpty()) {
                Document doc = new Document(para.trim());
                initialSegments.add(doc);
            }
        }

        // 2. 对每个段落再用 TokenTextSplitter 做 token 切分
        List<Document> finalChunks = new ArrayList<>();
        for (Document segment : initialSegments) {
            List<Document> tokenChunks = tokenSplitter.apply(List.of(segment));
            finalChunks.addAll(tokenChunks);
        }
        return finalChunks;
    }
}
